from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
import os

from demo02 import model

os.environ["DASHSCOPE_API_KEY"] = 'sk-5bc0688c5761427cadb9df012e589136'
os.environ["OPENAI_API_KEY"] = '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'
os.environ["TAVILY_API_KEY"] = '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'


# 构造文档
documents = [
    Document(
        page_content="狗是很好的伴侣，以忠诚和友善而闻名。",
        metadata={"source": "mammal-pets-doc"},
    ),
    Document(
        page_content="猫是独立的宠物，经常享受自己的空间。",
        metadata={"source": "mammal-pets-doc"},
    ),
    Document(
        page_content="金鱼是深受初学者欢迎的宠物，需要相对简单的护理。",
        metadata={"source": "fish-pets-doc"},
    ),
    Document(
        page_content="鹦鹉是一种聪明的鸟类，能够模仿人类的语言。",
        metadata={"source": "bird-pets-doc"},
    ),
    Document(
        page_content="兔子是群居动物，需要足够的空间来跳跃。",
        metadata={"source": "mammal-pets-doc"},
    ),
]


# 构造 prompt
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough

message = """
仅使用提供的上下文回答此问题。
{question}

上下文:
{context}
"""
prompt = ChatPromptTemplate.from_messages([("human", message)])


from langchain_chroma import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
# 创建 Chroma 向量存储
vectorstore = Chroma.from_documents(
    documents,
    embedding=DashScopeEmbeddings(),
)

# 查询向量存储
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 1},
)

# 使用 Tongyi LLM，并设置温度为 1，代表模型会更加随机，但也会更加不确定
from langchain_community.llms import Tongyi
llm = Tongyi(temperature=1)

# 构建 RAG 链
rag_chain = {"context": retriever, "question": RunnablePassthrough()} | prompt | llm

# 使用 RAG 链并打印结果
response = rag_chain.invoke("告诉我关于猫的事")
print(response)